from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from typing import Dict, Any
from ..db.database import get_db
from ..models.workflow import Workflow
from ..core.workflow_engine import WorkflowEngine
from ..core.ml_modules import (
    DataLoader, Preprocessor, ModelTrainer,
    FeatureSelector, ModelEvaluator, DataVisualizer
)
import pandas as pd

router = APIRouter()

@router.post("/workflows/{workflow_id}/execute")
async def execute_workflow(
    workflow_id: int,
    initial_data: Dict[str, Any] = None,
    db: Session = Depends(get_db)
):
    # 获取工作流配置
    workflow = db.query(Workflow).filter(Workflow.id == workflow_id).first()
    if not workflow:
        raise HTTPException(status_code=404, detail="Workflow not found")

    # 创建工作流引擎
    engine = WorkflowEngine()

    # 根据节点配置创建模块
    for node_id, node_data in workflow.nodes.items():
        module_type = node_data.get("type")
        if module_type == "data_loader":
            module = DataLoader()
        elif module_type == "preprocessor":
            module = Preprocessor()
        elif module_type == "model_trainer":
            module = ModelTrainer()
        elif module_type == "feature_selector":
            module = FeatureSelector()
        elif module_type == "model_evaluator":
            module = ModelEvaluator()
        elif module_type == "data_visualizer":
            module = DataVisualizer()
        else:
            continue

        # 设置模块参数
        module.set_params(node_data.get("params", {}))
        engine.add_module(node_id, module)

    # 设置执行顺序
    execution_order = []
    for edge in workflow.edges.values():
        if edge["source"] not in execution_order:
            execution_order.append(edge["source"])
        if edge["target"] not in execution_order:
            execution_order.append(edge["target"])
    engine.set_execution_order(execution_order)

    # 执行工作流
    try:
        results = engine.execute(pd.DataFrame(initial_data) if initial_data else None)
        visualization = engine.visualize_results()
        
        return {
            "status": "success",
            "results": results,
            "visualization": visualization
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e)) 